"""
增强版代码推荐系统
基于CodeQL数据流分析结果的智能代码推荐
"""

import os
import json
from typing import List, Dict, Optional, Tuple, Set
from neo4j import GraphDatabase, Driver
from dotenv import load_dotenv
from code_recommender import CodeRecommender


class EnhancedCodeRecommender(CodeRecommender):
    """基于数据流分析的增强代码推荐器"""
    
    def __init__(self, driver: Driver):
        super().__init__(driver)
        self.driver = driver
    
    def get_dataflow_based_recommendations(self, 
                                         current_function: str,
                                         data_type: str = "any",
                                         limit: int = 10) -> List[Dict]:
        """
        基于数据流分析的函数推荐
        
        Args:
            current_function: 当前函数名
            data_type: 数据类型 (user_input, file_data, network_data, etc.)
            limit: 返回数量限制
            
        Returns:
            基于数据流的推荐列表
        """
        with self.driver.session() as session:
            # 查询基于数据流模式的推荐
            dataflow_query = """
            MATCH (current:Function)-[:CALLS]->(intermediate:Function)-[:CALLS]->(recommended:Function)
            WHERE current.name CONTAINS $function_name 
               OR current.id CONTAINS $function_name
            WITH recommended, count(intermediate) as flow_paths
            ORDER BY flow_paths DESC
            LIMIT $limit
            RETURN recommended.id as function_id,
                   recommended.name as function_name,
                   recommended.file as function_file,
                   flow_paths,
                   'dataflow_pattern' as recommendation_type,
                   '基于数据流模式的函数调用链' as reason
            """
            
            results = []
            query_results = session.run(dataflow_query, 
                                      function_name=current_function,
                                      limit=limit)
            
            for record in query_results:
                results.append({
                    "function_id": record["function_id"],
                    "function_name": record["function_name"],
                    "function_file": record["function_file"],
                    "flow_paths": record["flow_paths"],
                    "recommendation_type": record["recommendation_type"],
                    "reason": record["reason"]
                })
            
            return results
    
    def get_security_aware_recommendations(self, 
                                         current_function: str,
                                         security_level: str = "safe",
                                         limit: int = 10) -> List[Dict]:
        """
        基于安全分析的函数推荐
        
        Args:
            current_function: 当前函数名
            security_level: 安全级别 (safe, warning, dangerous)
            limit: 返回数量限制
            
        Returns:
            安全感知的推荐列表
        """
        with self.driver.session() as session:
            # 查询安全相关的函数推荐
            security_query = """
            MATCH (current:Function)-[:CALLS]->(recommended:Function)
            WHERE current.name CONTAINS $function_name 
               OR current.id CONTAINS $function_name
            WITH recommended, 
                 size([(recommended)-[:CALLS]->(f:Function) WHERE f.name CONTAINS 'validate' OR f.name CONTAINS 'sanitize' | 1]) as security_calls
            ORDER BY security_calls DESC
            LIMIT $limit
            RETURN recommended.id as function_id,
                   recommended.name as function_name,
                   recommended.file as function_file,
                   security_calls,
                   'security_aware' as recommendation_type,
                   '基于安全最佳实践的推荐' as reason
            """
            
            results = []
            query_results = session.run(security_query, 
                                      function_name=current_function,
                                      limit=limit)
            
            for record in query_results:
                results.append({
                    "function_id": record["function_id"],
                    "function_name": record["function_name"],
                    "function_file": record["function_file"],
                    "security_calls": record["security_calls"],
                    "recommendation_type": record["recommendation_type"],
                    "reason": record["reason"]
                })
            
            return results
    
    def get_control_flow_recommendations(self, 
                                       current_function: str,
                                       flow_type: str = "conditional",
                                       limit: int = 10) -> List[Dict]:
        """
        基于控制流分析的函数推荐
        
        Args:
            current_function: 当前函数名
            flow_type: 控制流类型 (conditional, loop, exception)
            limit: 返回数量限制
            
        Returns:
            基于控制流的推荐列表
        """
        with self.driver.session() as session:
            # 查询控制流相关的函数推荐
            controlflow_query = """
            MATCH (current:Function)-[:CALLS]->(recommended:Function)
            WHERE current.name CONTAINS $function_name 
               OR current.id CONTAINS $function_name
            WITH recommended,
                 size([(recommended)-[:CALLS]->(f:Function) WHERE f.name CONTAINS 'check' OR f.name CONTAINS 'validate' | 1]) as control_calls
            ORDER BY control_calls DESC
            LIMIT $limit
            RETURN recommended.id as function_id,
                   recommended.name as function_name,
                   recommended.file as function_file,
                   control_calls,
                   'control_flow' as recommendation_type,
                   '基于控制流模式的推荐' as reason
            """
            
            results = []
            query_results = session.run(controlflow_query, 
                                      function_name=current_function,
                                      limit=limit)
            
            for record in query_results:
                results.append({
                    "function_id": record["function_id"],
                    "function_name": record["function_name"],
                    "function_file": record["function_file"],
                    "control_calls": record["control_calls"],
                    "recommendation_type": record["recommendation_type"],
                    "reason": record["reason"]
                })
            
            return results
    
    def get_vulnerability_aware_recommendations(self, 
                                              current_function: str,
                                              vulnerability_type: str = "any",
                                              limit: int = 10) -> List[Dict]:
        """
        基于漏洞分析的函数推荐
        
        Args:
            current_function: 当前函数名
            vulnerability_type: 漏洞类型 (buffer_overflow, sql_injection, etc.)
            limit: 返回数量限制
            
        Returns:
            漏洞感知的推荐列表
        """
        with self.driver.session() as session:
            # 查询安全函数推荐
            safe_functions_query = """
            MATCH (current:Function)-[:CALLS]->(recommended:Function)
            WHERE current.name CONTAINS $function_name 
               OR current.id CONTAINS $function_name
            WITH recommended,
                 size([(recommended)-[:CALLS]->(f:Function) WHERE 
                   f.name CONTAINS 'safe' OR 
                   f.name CONTAINS 'secure' OR 
                   f.name CONTAINS 'validate' OR
                   f.name CONTAINS 'sanitize' | 1]) as safety_score
            ORDER BY safety_score DESC
            LIMIT $limit
            RETURN recommended.id as function_id,
                   recommended.name as function_name,
                   recommended.file as function_file,
                   safety_score,
                   'vulnerability_aware' as recommendation_type,
                   '基于安全漏洞分析的推荐' as reason
            """
            
            results = []
            query_results = session.run(safe_functions_query, 
                                      function_name=current_function,
                                      limit=limit)
            
            for record in query_results:
                results.append({
                    "function_id": record["function_id"],
                    "function_name": record["function_name"],
                    "function_file": record["function_file"],
                    "safety_score": record["safety_score"],
                    "recommendation_type": record["recommendation_type"],
                    "reason": record["reason"]
                })
            
            return results
    
    def get_comprehensive_recommendations(self, 
                                        current_function: str,
                                        context: Dict = None,
                                        limit: int = 15) -> Dict[str, List[Dict]]:
        """
        获取综合推荐结果
        
        Args:
            current_function: 当前函数名
            context: 上下文信息 (文件路径、行号、数据类型等)
            limit: 每种推荐类型的数量限制
            
        Returns:
            包含多种推荐类型的综合结果
        """
        context = context or {}
        
        # 获取各种类型的推荐
        recommendations = {
            "basic": self.get_function_recommendations(current_function, limit=limit//3),
            "dataflow": self.get_dataflow_based_recommendations(current_function, limit=limit//3),
            "security": self.get_security_aware_recommendations(current_function, limit=limit//3),
            "control_flow": self.get_control_flow_recommendations(current_function, limit=limit//3),
            "vulnerability_aware": self.get_vulnerability_aware_recommendations(current_function, limit=limit//3)
        }
        
        # 如果有上下文信息，获取上下文感知推荐
        if context.get("current_file") and context.get("current_line"):
            recommendations["context_aware"] = self.get_context_aware_recommendations(
                context["current_file"], context["current_line"], limit=limit//3
            )
        
        return recommendations
    
    def analyze_code_patterns(self, 
                            function_name: str) -> Dict[str, any]:
        """
        分析代码模式
        
        Args:
            function_name: 函数名
            
        Returns:
            代码模式分析结果
        """
        with self.driver.session() as session:
            # 分析函数调用模式
            pattern_query = """
            MATCH (f:Function)
            WHERE f.name CONTAINS $function_name OR f.id CONTAINS $function_name
            OPTIONAL MATCH (f)-[:CALLS]->(callee:Function)
            OPTIONAL MATCH (caller:Function)-[:CALLS]->(f)
            RETURN f.name as function_name,
                   f.file as function_file,
                   count(callee) as outgoing_calls,
                   count(caller) as incoming_calls,
                   collect(DISTINCT callee.name)[0..5] as common_callees,
                   collect(DISTINCT caller.name)[0..5] as common_callers
            """
            
            result = session.run(pattern_query, function_name=function_name).single()
            
            if result:
                return {
                    "function_name": result["function_name"],
                    "function_file": result["function_file"],
                    "outgoing_calls": result["outgoing_calls"],
                    "incoming_calls": result["incoming_calls"],
                    "common_callees": result["common_callees"],
                    "common_callers": result["common_callers"],
                    "complexity_score": result["outgoing_calls"] + result["incoming_calls"]
                }
            
            return {}
    
    def get_smart_recommendations(self, 
                                current_function: str,
                                user_context: Dict = None) -> Dict[str, any]:
        """
        智能推荐：结合多种分析结果
        
        Args:
            current_function: 当前函数名
            user_context: 用户上下文 (编程经验、项目类型等)
            
        Returns:
            智能推荐结果
        """
        user_context = user_context or {}
        
        # 分析代码模式
        pattern_analysis = self.analyze_code_patterns(current_function)
        
        # 获取综合推荐
        recommendations = self.get_comprehensive_recommendations(current_function, user_context)
        
        # 计算推荐分数
        scored_recommendations = []
        for category, recs in recommendations.items():
            for rec in recs:
                score = self._calculate_recommendation_score(rec, pattern_analysis, user_context)
                rec["score"] = score
                rec["category"] = category
                scored_recommendations.append(rec)
        
        # 按分数排序
        scored_recommendations.sort(key=lambda x: x["score"], reverse=True)
        
        return {
            "pattern_analysis": pattern_analysis,
            "top_recommendations": scored_recommendations[:10],
            "recommendations_by_category": recommendations,
            "user_context": user_context
        }
    
    def _calculate_recommendation_score(self, 
                                      recommendation: Dict, 
                                      pattern_analysis: Dict, 
                                      user_context: Dict) -> float:
        """计算推荐分数"""
        score = 0.0
        
        # 基础分数
        if "call_frequency" in recommendation:
            score += recommendation["call_frequency"] * 0.3
        if "flow_paths" in recommendation:
            score += recommendation["flow_paths"] * 0.2
        if "security_calls" in recommendation:
            score += recommendation["security_calls"] * 0.15
        if "safety_score" in recommendation:
            score += recommendation["safety_score"] * 0.15
        
        # 模式匹配分数
        if pattern_analysis:
            if recommendation["function_name"] in pattern_analysis.get("common_callees", []):
                score += 0.1
            if recommendation["function_name"] in pattern_analysis.get("common_callers", []):
                score += 0.1
        
        # 用户上下文分数
        if user_context.get("prefer_security"):
            if recommendation["recommendation_type"] in ["security_aware", "vulnerability_aware"]:
                score += 0.2
        
        return score


def main():
    """演示增强版代码推荐功能"""
    from code_recommender import get_driver
    
    driver = get_driver()
    recommender = EnhancedCodeRecommender(driver)
    
    try:
        print("🚀 增强版代码推荐系统演示")
        print("=" * 60)
        
        # 演示智能推荐
        test_function = "main"
        user_context = {
            "prefer_security": True,
            "current_file": "src/main.cpp",
            "current_line": 10
        }
        
        print(f"\n🎯 智能推荐分析: '{test_function}'")
        smart_recs = recommender.get_smart_recommendations(test_function, user_context)
        
        print(f"\n📊 代码模式分析:")
        pattern = smart_recs["pattern_analysis"]
        if pattern:
            print(f"   函数名: {pattern['function_name']}")
            print(f"   文件: {pattern['function_file']}")
            print(f"   出站调用: {pattern['outgoing_calls']}")
            print(f"   入站调用: {pattern['incoming_calls']}")
            print(f"   复杂度分数: {pattern['complexity_score']}")
        
        print(f"\n🏆 顶级推荐 (按分数排序):")
        for i, rec in enumerate(smart_recs["top_recommendations"][:5], 1):
            print(f"  {i}. {rec['function_name']} (分数: {rec['score']:.2f})")
            print(f"     类型: {rec['recommendation_type']}")
            print(f"     理由: {rec['reason']}")
            print()
        
        # 演示数据流推荐
        print(f"\n🌊 数据流分析推荐:")
        dataflow_recs = recommender.get_dataflow_based_recommendations(test_function, limit=3)
        for i, rec in enumerate(dataflow_recs, 1):
            print(f"  {i}. {rec['function_name']} (流路径: {rec['flow_paths']})")
            print(f"     理由: {rec['reason']}")
            print()
        
        # 演示安全推荐
        print(f"\n🔒 安全感知推荐:")
        security_recs = recommender.get_security_aware_recommendations(test_function, limit=3)
        for i, rec in enumerate(security_recs, 1):
            print(f"  {i}. {rec['function_name']} (安全调用: {rec['security_calls']})")
            print(f"     理由: {rec['reason']}")
            print()
        
    finally:
        driver.close()


if __name__ == "__main__":
    main()
